from itertools import product
import pandas as pd
import numpy as np
import talib  # 用于技术指标计算
import base.stockbase as sb
from datetime import date


def get_bottom_divergence(symbol: str, period: str,
                          stock_data: pd.DataFrame, data_type: str, calculate_type: str):
    try:
        print(
            f"\n正在处理股票: {symbol} 关于数据类型为:{data_type}  计算值类型为: {calculate_type}的{period}数据")
        stock_data['MACD_DIF'], stock_data['MACD_DEA'], stock_data['MACD_HIST'] = talib.MACD(
            stock_data[data_type].to_numpy(), fastperiod=12, slowperiod=26, signalperiod=9)

        stock_data['MACD_HIST_DIFF'] = stock_data['MACD_HIST'].diff().bfill()

        # 识别价格波谷
        N = 26  # 周期参数
        stock_data['price_min'] = stock_data[data_type].rolling(
            window=N, min_periods=1).min()
        is_price_new_low = (stock_data[calculate_type]
                            <= stock_data['price_min'])  # 当前是否为新低

        # 识别指标波谷
        stock_data['macd_hist_min'] = stock_data['MACD_HIST'].rolling(
            window=N, min_periods=1).min()
        is_macd_not_new_low = (
            (stock_data['MACD_HIST'] > stock_data['macd_hist_min']) & (
                stock_data['MACD_HIST'] < 0)
            & (stock_data['MACD_HIST_DIFF'] > 0)
        )  # 指标未新低

        stock_data['bottom_divergence'] = np.where(
            (is_price_new_low) & (is_macd_not_new_low),  # 价格新低且指标未新低
            1,  # 信号触发
            0   # 无信号
        )

        stock_data = stock_data[stock_data['bottom_divergence'] == 1]
        stock_data = stock_data[stock_data['trade_date'] > date(2025, 1, 1)]
        if not stock_data.empty:
            return stock_data
        else:
            return pd.DataFrame()
    except Exception as e:
        print(f"计算 {symbol} 数据失败: {e}")


def main():
    data_cal_types = get_data_cal_types()
    all_daily_data = sb.get_stock_data('', 'daily')
    daily_symbols = all_daily_data['symbol'].unique()

    all_weekly_data = sb.get_stock_data('', 'weekly')
    weekly_symbols = all_weekly_data['symbol'].unique()
    for data_cal_type in data_cal_types:
        daily_dfs = []
        for daily_symbol in daily_symbols:
            symbol_daily_data = all_daily_data.loc[all_daily_data['symbol'] == daily_symbol].copy(
            )
            daily_dfs.append(get_bottom_divergence(
                daily_symbol, 'daily', symbol_daily_data, data_cal_type[0], data_cal_type[1]))

        weekly_dfs = []
        for weekly_symbol in weekly_symbols:
            symbol_weekly_data = all_weekly_data.loc[all_weekly_data['symbol'] == weekly_symbol].copy(
            )
            weekly_dfs.append(get_bottom_divergence(
                weekly_symbol, 'weekly', symbol_weekly_data, data_cal_type[0], data_cal_type[1]))

        str_today = date.today().strftime('%Y-%m-%d-%H-%M')
        csv_dir = 'output/indicator/macd/'
        pd.concat(daily_dfs).to_csv(csv_dir + '_'.join(['daily', str_today,
                                                        data_cal_type[0], data_cal_type[1]]) + '.csv')
        pd.concat(weekly_dfs).to_csv(csv_dir + '_'.join(['weekly', str_today,
                                                         data_cal_type[0], data_cal_type[1]]) + '.csv')


def get_data_cal_types():
    # 计算的数据类型
    data_cal_types = ['low', 'close']
    result = list(product(data_cal_types, repeat=2))
    return result


if __name__ == "__main__":
    main()
